/***************************************************************************
 *  Copyright (C) 2017 Codeplay Software Limited
 *  This Source Code Form is subject to the terms of the Mozilla
 *  Public License v. 2.0. If a copy of the MPL was not distributed
 *  with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 *
 *  SyclMemoryModel.h
 *
 *  Description:
 *    Interface for SYCL buffers to behave as a non-dereferenceable pointer
 *    Interface for Placeholder accessor to behave as a pointer on both host
 *    and device
 *
 * Authors:
 *
 *    Ruyman Reyes   Codeplay Software Ltd.
 *    Mehdi Goli     Codeplay Software Ltd.
 *    Vanya Yaneva   Codeplay Software Ltd.
 *
 **************************************************************************/

#if defined(EIGEN_USE_SYCL) && !defined(EIGEN_CXX11_TENSOR_TENSOR_SYCL_STORAGE_MEMORY_H)
#define EIGEN_CXX11_TENSOR_TENSOR_SYCL_STORAGE_MEMORY_H

#include <CL/sycl.hpp>
#ifdef EIGEN_EXCEPTIONS
#include <stdexcept>
#endif
#include <cstddef>
#include <queue>
#include <set>
#include <unordered_map>

namespace Eigen {
namespace TensorSycl {
    namespace internal {

        using sycl_acc_target = cl::sycl::access::target;
        using sycl_acc_mode = cl::sycl::access::mode;

        /**
 * Default values for template arguments
 */
        using buffer_data_type_t = uint8_t;
        const sycl_acc_target default_acc_target = sycl_acc_target::global_buffer;
        const sycl_acc_mode default_acc_mode = sycl_acc_mode::read_write;

        /**
 * PointerMapper
 *  Associates fake pointers with buffers.
 *
 */
        class PointerMapper
        {
        public:
            using base_ptr_t = std::intptr_t;

            /* Structure of a virtual pointer
   *
   * |================================================|
   * |               POINTER ADDRESS                  |
   * |================================================|
   */
            struct virtual_pointer_t
            {
                /* Type for the pointers
     */
                base_ptr_t m_contents;

                /** Conversions from virtual_pointer_t to
     * void * should just reinterpret_cast the integer number
     */
                operator void*() const { return reinterpret_cast<void*>(m_contents); }

                /**
     * Convert back to the integer number.
     */
                operator base_ptr_t() const { return m_contents; }

                /**
     * Add a certain value to the pointer to create a
     * new pointer to that offset
     */
                virtual_pointer_t operator+(size_t off) { return m_contents + off; }

                /* Numerical order for sorting pointers in containers. */
                bool operator<(virtual_pointer_t rhs) const { return (static_cast<base_ptr_t>(m_contents) < static_cast<base_ptr_t>(rhs.m_contents)); }

                bool operator>(virtual_pointer_t rhs) const { return (static_cast<base_ptr_t>(m_contents) > static_cast<base_ptr_t>(rhs.m_contents)); }

                /**
     * Numerical order for sorting pointers in containers
     */
                bool operator==(virtual_pointer_t rhs) const { return (static_cast<base_ptr_t>(m_contents) == static_cast<base_ptr_t>(rhs.m_contents)); }

                /**
     * Simple forward to the equality overload.
     */
                bool operator!=(virtual_pointer_t rhs) const { return !(this->operator==(rhs)); }

                /**
     * Converts a void * into a virtual pointer structure.
     * Note that this will only work if the void * was
     * already a virtual_pointer_t, but we have no way of
     * checking
     */
                virtual_pointer_t(const void* ptr) : m_contents(reinterpret_cast<base_ptr_t>(ptr)){};

                /**
     * Creates a virtual_pointer_t from the given integer
     * number
     */
                virtual_pointer_t(base_ptr_t u) : m_contents(u){};
            };

            /* Definition of a null pointer
   */
            const virtual_pointer_t null_virtual_ptr = nullptr;

            /**
   * Whether if a pointer is null or not.
   * A pointer is nullptr if the value is of null_virtual_ptr
   */
            static inline bool is_nullptr(virtual_pointer_t ptr) { return (static_cast<void*>(ptr) == nullptr); }

            /* basic type for all buffers
   */
            using buffer_t = cl::sycl::buffer_mem;

            /**
   * Node that stores information about a device allocation.
   * Nodes are sorted by size to organise a free list of nodes
   * that can be recovered.
   */
            struct pMapNode_t
            {
                buffer_t m_buffer;
                size_t m_size;
                bool m_free;

                pMapNode_t(buffer_t b, size_t size, bool f) : m_buffer{b}, m_size{size}, m_free{f} { m_buffer.set_final_data(nullptr); }

                bool operator<=(const pMapNode_t& rhs) { return (m_size <= rhs.m_size); }
            };

            /** Storage of the pointer / buffer tree
   */
            using pointerMap_t = std::map<virtual_pointer_t, pMapNode_t>;

            /**
   * Obtain the insertion point in the pointer map for
   * a pointer of the given size.
   * \param requiredSize Size attemted to reclaim
   */
            typename pointerMap_t::iterator get_insertion_point(size_t requiredSize)
            {
                typename pointerMap_t::iterator retVal;
                bool reuse = false;
                if (!m_freeList.empty())
                {
                    // try to re-use an existing block
                    for (auto freeElem : m_freeList)
                    {
                        if (freeElem->second.m_size >= requiredSize)
                        {
                            retVal = freeElem;
                            reuse = true;
                            // Element is not going to be free anymore
                            m_freeList.erase(freeElem);
                            break;
                        }
                    }
                }
                if (!reuse)
                {
                    retVal = std::prev(m_pointerMap.end());
                }
                return retVal;
            }

            /**
   * Returns an iterator to the node that stores the information
   * of the given virtual pointer from the given pointer map structure.
   * If pointer is not found, throws std::out_of_range.
   * If the pointer map structure is empty, throws std::out_of_range
   *
   * \param pMap the pointerMap_t structure storing all the pointers
   * \param virtual_pointer_ptr The virtual pointer to obtain the node of
   * \throws std::out:of_range if the pointer is not found or pMap is empty
   */
            typename pointerMap_t::iterator get_node(const virtual_pointer_t ptr)
            {
                if (this->count() == 0)
                {
                    m_pointerMap.clear();
                    EIGEN_THROW_X(std::out_of_range("There are no pointers allocated\n"));
                }
                if (is_nullptr(ptr))
                {
                    m_pointerMap.clear();
                    EIGEN_THROW_X(std::out_of_range("Cannot access null pointer\n"));
                }
                // The previous element to the lower bound is the node that
                // holds this memory address
                auto node = m_pointerMap.lower_bound(ptr);
                // If the value of the pointer is not the one of the node
                // then we return the previous one
                if (node == std::end(m_pointerMap))
                {
                    --node;
                }
                else if (node->first != ptr)
                {
                    if (node == std::begin(m_pointerMap))
                    {
                        m_pointerMap.clear();
                        EIGEN_THROW_X(std::out_of_range("The pointer is not registered in the map\n"));
                    }
                    --node;
                }

                return node;
            }

            /* get_buffer.
   * Returns a buffer from the map using the pointer address
   */
            template <typename buffer_data_type = buffer_data_type_t> cl::sycl::buffer<buffer_data_type, 1> get_buffer(const virtual_pointer_t ptr)
            {
                using sycl_buffer_t = cl::sycl::buffer<buffer_data_type, 1>;

                // get_node() returns a `buffer_mem`, so we need to cast it to a `buffer<>`.
                // We can do this without the `buffer_mem` being a pointer, as we
                // only declare member variables in the base class (`buffer_mem`) and not in
                // the child class (`buffer<>).
                auto node = get_node(ptr);
                eigen_assert(node->first == ptr || node->first < ptr);
                eigen_assert(ptr < static_cast<virtual_pointer_t>(node->second.m_size + node->first));
                return *(static_cast<sycl_buffer_t*>(&node->second.m_buffer));
            }

            /**
   * @brief Returns an accessor to the buffer of the given virtual pointer
   * @param accessMode
   * @param accessTarget
   * @param ptr The virtual pointer
   */
            template <sycl_acc_mode access_mode = default_acc_mode,
                      sycl_acc_target access_target = default_acc_target,
                      typename buffer_data_type = buffer_data_type_t>
            cl::sycl::accessor<buffer_data_type, 1, access_mode, access_target> get_access(const virtual_pointer_t ptr)
            {
                auto buf = get_buffer<buffer_data_type>(ptr);
                return buf.template get_access<access_mode, access_target>();
            }

            /**
   * @brief Returns an accessor to the buffer of the given virtual pointer
   *        in the given command group scope
   * @param accessMode
   * @param accessTarget
   * @param ptr The virtual pointer
   * @param cgh Reference to the command group scope
   */
            template <sycl_acc_mode access_mode = default_acc_mode,
                      sycl_acc_target access_target = default_acc_target,
                      typename buffer_data_type = buffer_data_type_t>
            cl::sycl::accessor<buffer_data_type, 1, access_mode, access_target> get_access(const virtual_pointer_t ptr, cl::sycl::handler& cgh)
            {
                auto buf = get_buffer<buffer_data_type>(ptr);
                return buf.template get_access<access_mode, access_target>(cgh);
            }

            /*
   * Returns the offset from the base address of this pointer.
   */
            inline std::ptrdiff_t get_offset(const virtual_pointer_t ptr)
            {
                // The previous element to the lower bound is the node that
                // holds this memory address
                auto node = get_node(ptr);
                auto start = node->first;
                eigen_assert(start == ptr || start < ptr);
                eigen_assert(ptr < start + node->second.m_size);
                return (ptr - start);
            }

            /*
   * Returns the number of elements by which the given pointer is offset from
   * the base address.
   */
            template <typename buffer_data_type> inline size_t get_element_offset(const virtual_pointer_t ptr)
            {
                return get_offset(ptr) / sizeof(buffer_data_type);
            }

            /**
   * Constructs the PointerMapper structure.
   */
            PointerMapper(base_ptr_t baseAddress = 4096) : m_pointerMap{}, m_freeList{}, m_baseAddress{baseAddress}
            {
                if (m_baseAddress == 0)
                {
                    EIGEN_THROW_X(std::invalid_argument("Base address cannot be zero\n"));
                }
            };

            /**
   * PointerMapper cannot be copied or moved
   */
            PointerMapper(const PointerMapper&) = delete;

            /**
   * Empty the pointer list
   */
            inline void clear()
            {
                m_freeList.clear();
                m_pointerMap.clear();
            }

            /* add_pointer.
   * Adds an existing pointer to the map and returns the virtual pointer id.
   */
            inline virtual_pointer_t add_pointer(const buffer_t& b) { return add_pointer_impl(b); }

            /* add_pointer.
   * Adds a pointer to the map and returns the virtual pointer id.
   */
            inline virtual_pointer_t add_pointer(buffer_t&& b) { return add_pointer_impl(b); }

            /**
   * @brief Fuses the given node with the previous nodes in the
   *        pointer map if they are free
   *
   * @param node A reference to the free node to be fused
   */
            void fuse_forward(typename pointerMap_t::iterator& node)
            {
                while (node != std::prev(m_pointerMap.end()))
                {
                    // if following node is free
                    // remove it and extend the current node with its size
                    auto fwd_node = std::next(node);
                    if (!fwd_node->second.m_free)
                    {
                        break;
                    }
                    auto fwd_size = fwd_node->second.m_size;
                    m_freeList.erase(fwd_node);
                    m_pointerMap.erase(fwd_node);

                    node->second.m_size += fwd_size;
                }
            }

            /**
   * @brief Fuses the given node with the following nodes in the
   *        pointer map if they are free
   *
   * @param node A reference to the free node to be fused
   */
            void fuse_backward(typename pointerMap_t::iterator& node)
            {
                while (node != m_pointerMap.begin())
                {
                    // if previous node is free, extend it
                    // with the size of the current one
                    auto prev_node = std::prev(node);
                    if (!prev_node->second.m_free)
                    {
                        break;
                    }
                    prev_node->second.m_size += node->second.m_size;

                    // remove the current node
                    m_freeList.erase(node);
                    m_pointerMap.erase(node);

                    // point to the previous node
                    node = prev_node;
                }
            }

            /* remove_pointer.
   * Removes the given pointer from the map.
   * The pointer is allowed to be reused only if ReUse if true.
   */
            template <bool ReUse = true> void remove_pointer(const virtual_pointer_t ptr)
            {
                if (is_nullptr(ptr))
                {
                    return;
                }
                auto node = this->get_node(ptr);

                node->second.m_free = true;
                m_freeList.emplace(node);

                // Fuse the node
                // with free nodes before and after it
                fuse_forward(node);
                fuse_backward(node);

                // If after fusing the node is the last one
                // simply remove it (since it is free)
                if (node == std::prev(m_pointerMap.end()))
                {
                    m_freeList.erase(node);
                    m_pointerMap.erase(node);
                }
            }

            /* count.
   * Return the number of active pointers (i.e, pointers that
   * have been malloc but not freed).
   */
            size_t count() const { return (m_pointerMap.size() - m_freeList.size()); }

        private:
            /* add_pointer_impl.
   * Adds a pointer to the map and returns the virtual pointer id.
   * BufferT is either a const buffer_t& or a buffer_t&&.
   */
            template <class BufferT> virtual_pointer_t add_pointer_impl(BufferT b)
            {
                virtual_pointer_t retVal = nullptr;
                size_t bufSize = b.get_count();
                pMapNode_t p{b, bufSize, false};
                // If this is the first pointer:
                if (m_pointerMap.empty())
                {
                    virtual_pointer_t initialVal{m_baseAddress};
                    m_pointerMap.emplace(initialVal, p);
                    return initialVal;
                }

                auto lastElemIter = get_insertion_point(bufSize);
                // We are recovering an existing free node
                if (lastElemIter->second.m_free)
                {
                    lastElemIter->second.m_buffer = b;
                    lastElemIter->second.m_free = false;

                    // If the recovered node is bigger than the inserted one
                    // add a new free node with the remaining space
                    if (lastElemIter->second.m_size > bufSize)
                    {
                        // create a new node with the remaining space
                        auto remainingSize = lastElemIter->second.m_size - bufSize;
                        pMapNode_t p2{b, remainingSize, true};

                        // update size of the current node
                        lastElemIter->second.m_size = bufSize;

                        // add the new free node
                        auto newFreePtr = lastElemIter->first + bufSize;
                        auto freeNode = m_pointerMap.emplace(newFreePtr, p2).first;
                        m_freeList.emplace(freeNode);
                    }

                    retVal = lastElemIter->first;
                }
                else
                {
                    size_t lastSize = lastElemIter->second.m_size;
                    retVal = lastElemIter->first + lastSize;
                    m_pointerMap.emplace(retVal, p);
                }
                return retVal;
            }

            /**
   * Compare two iterators to pointer map entries according to
   * the size of the allocation on the device.
   */
            struct SortBySize
            {
                bool operator()(typename pointerMap_t::iterator a, typename pointerMap_t::iterator b) const
                {
                    return ((a->first < b->first) && (a->second <= b->second)) || ((a->first < b->first) && (b->second <= a->second));
                }
            };

            /* Maps the pointer addresses to buffer and size pairs.
   */
            pointerMap_t m_pointerMap;

            /* List of free nodes available for re-using
   */
            std::set<typename pointerMap_t::iterator, SortBySize> m_freeList;

            /* Base address used when issuing the first virtual pointer, allows users
   * to specify alignment. Cannot be zero. */
            std::intptr_t m_baseAddress;
        };

        /* remove_pointer.
 * Removes the given pointer from the map.
 * The pointer is allowed to be reused only if ReUse if true.
 */
        template <> inline void PointerMapper::remove_pointer<false>(const virtual_pointer_t ptr)
        {
            if (is_nullptr(ptr))
            {
                return;
            }
            m_pointerMap.erase(this->get_node(ptr));
        }

        /**
 * Malloc-like interface to the pointer-mapper.
 * Given a size, creates a byte-typed buffer and returns a
 * fake pointer to keep track of it.
 * \param size Size in bytes of the desired allocation
 * \throw cl::sycl::exception if error while creating the buffer
 */
        inline void* SYCLmalloc(size_t size, PointerMapper& pMap)
        {
            if (size == 0)
            {
                return nullptr;
            }
            // Create a generic buffer of the given size
            using buffer_t = cl::sycl::buffer<buffer_data_type_t, 1>;
            auto thePointer = pMap.add_pointer(buffer_t(cl::sycl::range<1>{size}));
            // Store the buffer on the global list
            return static_cast<void*>(thePointer);
        }

        /**
 * Free-like interface to the pointer mapper.
 * Given a fake-pointer created with the virtual-pointer malloc,
 * destroys the buffer and remove it from the list.
 * If ReUse is false, the pointer is not added to the freeList,
 * it should be false only for sub-buffers.
 */
        template <bool ReUse = true, typename PointerMapper> inline void SYCLfree(void* ptr, PointerMapper& pMap) { pMap.template remove_pointer<ReUse>(ptr); }

        /**
 * Clear all the memory allocated by SYCL.
 */
        template <typename PointerMapper> inline void SYCLfreeAll(PointerMapper& pMap) { pMap.clear(); }

        template <cl::sycl::access::mode AcMd, typename T> struct RangeAccess
        {
            static const auto global_access = cl::sycl::access::target::global_buffer;
            static const auto is_place_holder = cl::sycl::access::placeholder::true_t;
            typedef T scalar_t;
            typedef scalar_t& ref_t;
            typedef typename cl::sycl::global_ptr<scalar_t>::pointer_t ptr_t;

            // the accessor type does not necessarily the same as T
            typedef cl::sycl::accessor<scalar_t, 1, AcMd, global_access, is_place_holder> accessor;

            typedef RangeAccess<AcMd, T> self_t;
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE RangeAccess(accessor access, size_t offset, std::intptr_t virtual_ptr)
                : access_(access), offset_(offset), virtual_ptr_(virtual_ptr)
            {
            }

            RangeAccess(cl::sycl::buffer<scalar_t, 1> buff = cl::sycl::buffer<scalar_t, 1>(cl::sycl::range<1>(1)))
                : access_{accessor{buff}}, offset_(0), virtual_ptr_(-1)
            {
            }

            // This should be only used for null constructor on the host side
            RangeAccess(std::nullptr_t) : RangeAccess() {}
            // This template parameter must be removed and scalar_t should be replaced
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ptr_t get_pointer() const { return (access_.get_pointer().get() + offset_); }
            template <typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE self_t& operator+=(Index offset)
            {
                offset_ += (offset);
                return *this;
            }
            template <typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE self_t operator+(Index offset) const
            {
                return self_t(access_, offset_ + offset, virtual_ptr_);
            }
            template <typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE self_t operator-(Index offset) const
            {
                return self_t(access_, offset_ - offset, virtual_ptr_);
            }
            template <typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE self_t& operator-=(Index offset)
            {
                offset_ -= offset;
                return *this;
            }

            // THIS IS FOR NULL COMPARISON ONLY
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE friend bool operator==(const RangeAccess& lhs, std::nullptr_t) { return ((lhs.virtual_ptr_ == -1)); }
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE friend bool operator!=(const RangeAccess& lhs, std::nullptr_t i) { return !(lhs == i); }

            // THIS IS FOR NULL COMPARISON ONLY
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE friend bool operator==(std::nullptr_t, const RangeAccess& rhs) { return ((rhs.virtual_ptr_ == -1)); }
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE friend bool operator!=(std::nullptr_t i, const RangeAccess& rhs) { return !(i == rhs); }
            // Prefix operator (Increment and return value)
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE self_t& operator++()
            {
                offset_++;
                return (*this);
            }

            // Postfix operator (Return value and increment)
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE self_t operator++(int i)
            {
                EIGEN_UNUSED_VARIABLE(i);
                self_t temp_iterator(*this);
                offset_++;
                return temp_iterator;
            }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::ptrdiff_t get_size() const { return (access_.get_count() - offset_); }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::ptrdiff_t get_offset() const { return offset_; }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void set_offset(std::ptrdiff_t offset) { offset_ = offset; }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ref_t operator*() const { return *get_pointer(); }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ref_t operator*() { return *get_pointer(); }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ptr_t operator->() = delete;

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ref_t operator[](int x) { return *(get_pointer() + x); }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ref_t operator[](int x) const { return *(get_pointer() + x); }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE scalar_t* get_virtual_pointer() const
            {
                return reinterpret_cast<scalar_t*>(virtual_ptr_ + (offset_ * sizeof(scalar_t)));
            }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE explicit operator bool() const { return (virtual_ptr_ != -1); }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE operator RangeAccess<AcMd, const T>() { return RangeAccess<AcMd, const T>(access_, offset_, virtual_ptr_); }

            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE operator RangeAccess<AcMd, const T>() const
            {
                return RangeAccess<AcMd, const T>(access_, offset_, virtual_ptr_);
            }
            // binding placeholder accessors to a command group handler for SYCL
            EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler& cgh) const { cgh.require(access_); }

        private:
            accessor access_;
            size_t offset_;
            std::intptr_t virtual_ptr_;  // the location of the buffer in the map
        };

        template <cl::sycl::access::mode AcMd, typename T> struct RangeAccess<AcMd, const T> : RangeAccess<AcMd, T>
        {
            typedef RangeAccess<AcMd, T> Base;
            using Base::Base;
        };

    }  // namespace internal
}  // namespace TensorSycl
}  // namespace Eigen

#endif  // EIGEN_CXX11_TENSOR_TENSOR_SYCL_STORAGE_MEMORY_H
